import torch
from torch import nn
import torchvision
from torch.nn import Conv2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid

'''
加载CIFAR-10数据集，定义一个简单的卷积神经网络模型，并使用TensorBoard记录输入和输出的图像张量。具体逻辑如下：
1. 加载CIFAR-10数据集并创建DataLoader。
2. 定义一个包含单层卷积的神经网络模型。
3. 遍历数据集，将输入图像通过模型进行前向传播，打印输入和输出张量的形状。
4. 使用make_grid合并图像张量，并将结果写入TensorBoard日志
'''
dataset = torchvision.datasets.CIFAR10("../../data", train=False,
                                       transform=torchvision.transforms.ToTensor(), download=True)

dataLoader = DataLoader(dataset,batch_size=64)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = Conv2d(in_channels=3,out_channels=6,kernel_size=3,stride=1,padding=0)

    def forward(self,x):
        x = self.conv1(x)
        return x

writer = SummaryWriter("../logs_conv2")
step = 0
model = Model()
for data in dataLoader:
    imgs,targets = data
    print(imgs.shape)

    output = model(imgs)
    print(output.shape)

    # 使用 make_grid 合并一批图像
    grid = make_grid(imgs)

    # 只使用输出张量的前 3 个通道
    output_grid = make_grid(output[:, :3, :, :])

    writer.add_image("input", grid, step)
    writer.add_image("output", output_grid, step)
    step += 1

writer.close()